#!/usr/bin/python3
# coding=utf-8
#
# Copyright (C) 2023-2024. Huawei Technologies Co., Ltd. All rights reserved.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# ===============================================================================

import numpy as np
import torch
import torch.nn.functional as F
import os
import struct

def col2im_torch(col, output_size, kernel_size, stride=1, padding=0):
    """
    PyTorch实现的col2im，使用F.fold
    col: [N, C*kH*kW, L] tensor
    output_size: (H, W) 输出图像的高和宽
    return: [N, C, H, W] tensor
    """
    if isinstance(kernel_size, int):
        kernel_size = (kernel_size, kernel_size)
    
    # 使用PyTorch的fold函数
    img = F.fold(col, output_size, kernel_size, stride=stride, padding=padding)
    return img

def gen_golden_data_col2im_torch():
    # 原始图像尺寸
    output_shape = (1, 1, 8, 256)  # [N, C, H, W]
    N, C, H, W = output_shape
    
    # col2im参数
    kernel_h = 2
    kernel_w = 2
    kernel_size = (kernel_h, kernel_w)
    stride_val = 1           # 步长
    padding_val = 0          # 填充
    
    # 计算col矩阵的尺寸
    out_H = (H + 2 * padding_val - kernel_h) // stride_val + 1
    out_W = (W + 2 * padding_val - kernel_w) // stride_val + 1
    L = out_H * out_W
    input_channels = C * kernel_h * kernel_w
    
    # 生成随机col矩阵作为输入
    col_shape = (N, input_channels, L)
    input_data_np = np.random.uniform(low=-1.0, high=1.0, size=col_shape).astype(np.float32)
    input_data = torch.from_numpy(input_data_np)
    
    # 计算col2im
    golden_tensor = col2im_torch(input_data, (H, W), kernel_size, stride_val, padding_val)
    golden = golden_tensor.numpy().astype(np.float32)
    
    # 创建目录
    os.makedirs("input", exist_ok=True)
    os.makedirs("output", exist_ok=True)
    
    # 根据 Col2ImCustomTilingData 结构体生成tiling数据
    # struct Col2ImCustomTilingData {
    #     uint32_t totalLength;  // N * input_channels * L (col矩阵大小)
    #     uint32_t tileNum;
    # 
    #     // 参数设置
    #     int32_t N;
    #     int32_t C;
    #     int32_t H;
    #     int32_t W;
    # 
    #     // 填充参数
    #     int32_t kernel_h;
    #     int32_t kernel_w;
    #     int32_t stride_val;
    #     int32_t padding_val;
    # };
    
    # 获取col矩阵的总长度
    total_length = N * input_channels * L  # totalLength (col矩阵大小)
    tile_num = 8  # 假设使用8个tile
    
    # 使用struct.pack确保正确的内存布局
    # 'I' = unsigned int (uint32_t), 'i' = int (int32_t)
    tiling_data = struct.pack('IIiiiiiiii', 
                             total_length,  # uint32_t totalLength
                             tile_num,      # uint32_t tileNum
                             N,             # int32_t N
                             C,             # int32_t C  
                             H,             # int32_t H
                             W,             # int32_t W
                             kernel_h,      # int32_t kernel_h
                             kernel_w,      # int32_t kernel_w
                             stride_val,    # int32_t stride_val
                             padding_val)   # int32_t padding_val
    
    # 保存tiling数据
    with open("input/input_tiling.bin", "wb") as f:
        f.write(tiling_data)
    
    # 保存输入和输出数据
    input_data_np.tofile("input/input_col.bin")
    golden.tofile("output/golden.bin")

if __name__ == "__main__":
    # 使用PyTorch实现
    gen_golden_data_col2im_torch()